import nsfw as n2
import os
import shutil
from nsfw._utils import extract_first_non_zero_digit


def test_predict_image() -> None:
    print()
    print(n2.predict_image("test_image_2.jpg"))


def test_read_images() -> None:
    source_directory = 'D:\project\com.weinian\image'
    # 读取指定目录
    for filename in os.listdir(source_directory):
        source_file = os.path.join(source_directory, filename)
        if n2.predict_image(source_file) > 0.5:
            destination_file = os.path.join("D:\\nsfw", filename)
            shutil.move(source_file, destination_file)


def test_read_images_group() -> None:
    source_directory = 'D:\project\com.weinian\image'
    destination_directory = "D:\smon\leve"
    # 读取指定目录
    for filename in os.listdir(source_directory):
        source_file = os.path.join(source_directory, filename)
        chance: float = n2.predict_image(source_file)
        print(chance)
        digit = str(extract_first_non_zero_digit(chance))
        path = destination_directory + digit
        destination_file = os.path.join(path, filename)
        shutil.move(source_file, destination_file)
